from PIL import Image
import numpy as np
import tensorflow as tf

class DigitRecognizer:
    def __init__(self, model_path='mnist_cnn.h5'):
        self.model = tf.keras.models.load_model(model_path)
    
    def preprocess_image(self, image_path):
        """处理输入图片"""
        img = Image.open(image_path)
        
        # 转换为灰度图
        img = img.convert('L')
        
        # 调整尺寸并保持数字居中
        img = img.resize((28, 28), resample=Image.BILINEAR)
        
        # 反转颜色（MNIST是白底黑字）
        img = Image.eval(img, lambda x: 255 - x)
        
        # 转换为numpy数组
        img_array = np.array(img).astype('float32')
        
        # 归一化并添加维度
        img_array = (img_array / 255.0).reshape(1, 28, 28, 1)
        return img_array
    
    def predict(self, image_path):
        """预测数字"""
        processed_img = self.preprocess_image(image_path)
        predictions = self.model.predict(processed_img)
        return np.argmax(predictions)

# 使用示例
if __name__ == "__main__":
    recognizer = DigitRecognizer()
    
    # 测试图片路径（需替换为实际路径）
    test_image = "2.png"
    
    # 执行预测
    prediction = recognizer.predict(test_image)
    print(f"Predicted digit: {prediction}")